import cv2
import argparse

import numpy as np


class yolov5():
    def __init__(self, model_path, confThreshold=0.8, nmsThreshold= 0.4, objThreshold = 0.8):
        with open("class.names", 'rt') as f:
            #f.read() 会将文件的内容一次性读取成一个字符串。
            #rstrip('\n') 去除字符串末尾的换行符（防止最后一行为空行）。
            #split('\n') 根据换行符 \n 将整个字符串分割成一个列表，每一行内容作为列表中的一个元素。这个操作假设文件中的每一行都代表一个类名。
            self.classes = f.read().rstrip('\n').split('\n')
            print(self.classes)
        #有多少种类别
        self.num_classes = len(self.classes)
        #判断字符串是否以指定字符或子字符串结尾
        if model_path.endswith('6.onnx'):
            self.inpHeight, self.inpWidth = 1280, 1280
            anchors = [[19, 27, 44, 40, 38, 94], [96, 68, 86, 152, 180, 137], [140, 301, 303, 264, 238, 542], [436, 615, 739, 380, 925, 792]]
            self.strides = np.array([8., 16., 32., 64.])
        else:
            self.inpHeight, self.inpWidth = 640, 640
            anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
            self.stride = np.array([8., 16., 32.])
        #`self.nl` 保存了锚框的层数或组数，这个变量可能用于定义网络中的层级结构（如输出特征图的层数）
        self.nl = len(anchors)
        print(self.nl)

        self.na = len(anchors[0]) // 2
        print('na', self.na)




if __name__ == '__main__':
    #初始化arrgparse结构
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgpath', type =str , default='images/zidane.jpg' ,help= 'image path')
    parser.add_argument('--modelpath', type= str, default='yolov5n.onnx', help= 'model path')
    #置信阈值：在机器学习和统计学中，用于确定预测结果是否可靠的阈值。当预测结果的置信度高于这个阈值时，认为预测结果是可靠的
    parser.add_argument('--confThreshold', type = float, default= 0.4, help= 'confidence threshold')
    parser.add_argument('--nmsThreshold', type= float, default= 0.4, help= 'nms threshold')
    # 【IOU门限】标记与真实框有交集的锚框的objectness（框内是否有物体）为1还是-1的分界线
    parser.add_argument('--objThreshold', type= float, default= 0.4, help= 'object threshold')
    #实例化parser结构体
    args = parser.parse_args()
    print(args)
    #实例化对象,实例化得时候会执行一次
    yolonet = yolov5(model_path= args.modelpath)

    img = cv2.imread(args.imgpath)
    cv2.imshow('img', img)
    cv2.waitKey(0)

